from data_visualization.data_visualization import DataVisualization


class SimplePltChart(DataVisualization):
    chart_type = ""
    y_label = ""
    y_values = []

    def __init__(self):
        DataVisualization.__init__(self)
        self.algorithm_name = "简单Matplotlib图表"
        self.ipynb_template_name = "simple_plt_chart-template.ipynb"
        
    def setYLabel(self, y_label):
        self.y_label = y_label
        self.y_values = self.data_source[self.y_label].values
        
    def implent(self): 
        DataVisualization.implent(self)
        ax = self.figure.add_subplot(111)
        ax.set_title(self.chart_title)
        if self.chart_type != "pie":
            ax.set_xlabel(self.x_label)
            ax.set_ylabel(self.y_label)  
        
        self.code = self.code + 'plt.title("#chart_title#")\n'
        if self.chart_type == "bar":
            ax.bar(self.x_values, self.y_values)
            self.code = self.code + "plt.bar(x_values, y_values)\n"
        elif self.chart_type == "line":
            ax.plot(self.x_values, self.y_values)
            self.code = self.code + "plt.plot(x_values, y_values)\n"
        elif self.chart_type == "scatter":
            ax.scatter(self.x_values, self.y_values)
            self.code = self.code + "plt.scatter(x_values, y_values)\n"
        elif self.chart_type == "pie":
            ax.axis('equal')
            ax.pie(x=self.y_values, labels=self.x_values, autopct='%1.1f%%')
            self.code = self.code + """
plt.axis('equal')
plt.pie(x=y_values, labels=x_values, autopct='%1.1f%%')
"""
            
        self.code = self.code + "plt.show()\n"
        self.saveToPng()

    def prepareIpynbItems(self):
        DataVisualization.prepareIpynbItems(self)
        self.ipynb_items["#chart_type#"] = self.chart_type
        self.ipynb_items["#y_label#"] = self.y_label
